#!/usr/bin/env python3
"""
PyTorch自定义算子构建脚本
"""

import argparse
import os
import subprocess
import sys


def build_with_setup():
    """使用setup.py构建自定义算子"""
    print("使用setup.py构建PyTorch自定义算子...")

    # 切换到当前目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(current_dir)

    # 运行setup.py构建
    result = subprocess.run([sys.executable, "setup.py", "build_ext", "--inplace"], capture_output=True, text=True)

    if result.returncode == 0:
        print("构建成功！")
        print("输出:", result.stdout)

        # 检查生成的文件
        lib_files = [f for f in os.listdir(current_dir) if f.startswith("dose_stat_torch") and f.endswith(".so")]
        if lib_files:
            print(f"生成的库文件: {lib_files}")
        else:
            print("警告: 未找到生成的库文件")

    else:
        print("构建失败！")
        print("错误:", result.stderr)
        return False

    return True


def build_with_cmake():
    """使用CMake构建自定义算子"""
    print("使用CMake构建PyTorch自定义算子...")

    current_dir = os.path.dirname(os.path.abspath(__file__))
    build_dir = os.path.join(current_dir, "build")

    # 创建构建目录
    os.makedirs(build_dir, exist_ok=True)
    os.chdir(build_dir)

    # 运行CMake配置
    cmake_result = subprocess.run(["cmake", "-DCMAKE_PREFIX_PATH=" + os.path.dirname(os.path.dirname(os.path.dirname(current_dir))), "-DCMAKE_BUILD_TYPE=Release", ".."], capture_output=True, text=True)

    if cmake_result.returncode != 0:
        print("CMake配置失败！")
        print("错误:", cmake_result.stderr)
        return False

    # 运行构建
    build_result = subprocess.run(["cmake", "--build", ".", "--config", "Release"], capture_output=True, text=True)

    if build_result.returncode == 0:
        print("CMake构建成功！")

        # 复制库文件到源代码目录
        lib_files = [f for f in os.listdir(build_dir) if f.startswith("dose_stat_torch") and f.endswith(".so")]
        for lib_file in lib_files:
            src_path = os.path.join(build_dir, lib_file)
            dst_path = os.path.join(current_dir, lib_file)

            import shutil

            shutil.copy2(src_path, dst_path)
            print(f"复制库文件: {lib_file} -> {dst_path}")

    else:
        print("CMake构建失败！")
        print("错误:", build_result.stderr)
        return False

    return True


def test_import():
    """测试导入自定义算子"""
    print("测试导入自定义算子...")

    current_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, current_dir)

    try:
        import dose_torch

        print("导入成功！")

        # 测试基本功能
        import torch

        # 创建测试数据
        voxel_volume = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
        voxel_dose = torch.tensor([10.0, 20.0, 30.0], dtype=torch.float64)

        # 测试accumulate函数
        binned_dose, accum_vol = dose_torch.accumulate(voxel_volume, voxel_dose)
        print(f"测试accumulate成功: binned_dose.shape={binned_dose.shape}, accum_vol.shape={accum_vol.shape}")

        return True

    except Exception as e:
        print(f"导入测试失败: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="PyTorch自定义算子构建脚本")
    parser.add_argument("--method", choices=["setup", "cmake", "both"], default="setup", help="构建方法 (setup.py 或 CMake)")
    parser.add_argument("--test", action="store_true", help="构建后测试导入")

    args = parser.parse_args()

    success = False

    if args.method in ["setup", "both"]:
        success = build_with_setup()
        if not success and args.method == "setup":
            sys.exit(1)

    if args.method in ["cmake", "both"] and (args.method == "both" or not success):
        success = build_with_cmake()
        if not success:
            sys.exit(1)

    if args.test and success:
        test_import()

    print("\n构建完成！")
    print("使用方法:")
    print("1. 在Python中导入: from dose_stat_torch import DoseStatTorch")
    print("2. 使用示例: binned_dose, accum_vol = DoseStatTorch.accumulate(voxel_volume, voxel_dose)")


if __name__ == "__main__":
    main()
